import datetime
import os

import pymongo
import pymongo.errors
import yaml
from loguru import logger


class MongoConnector:

    def __init__(self):
        self.db_config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.yml")
        config_content = self.read_config()
        self.db_addr = config_content["mongodb"]["addr"]
        self.db_port = config_content["mongodb"]["port"]
        self.db_username = config_content["mongodb"]["username"]
        self.db_password = config_content["mongodb"]["password"]

        self.db_client: pymongo.MongoClient = None
        self.selected_db = None
        self.selected_collection: pymongo.collection.Collection = None

    def read_config(self):
        config_file = open(self.db_config_path, 'r', encoding='utf-8')
        config_content = yaml.load(config_file, Loader=yaml.FullLoader)
        # logger.debug("read config file: {}", self.db_config_path)
        config_file.close()
        return config_content

    def init_client(self):
        self.db_client = pymongo.MongoClient(self.db_addr, self.db_port)
        if self.db_username != "":
            mongo_auth = self.db_client.admin  # 或 mongo_client['admin'] admin为authenticationDatabase
            mongo_auth.authenticate(self.db_username, self.db_password)
        # logger.info(self.db_client.server_info())

    def close_client(self):
        self.db_client.close()
        logger.info("Mongodb Closed")

    def select_db(self, database: str):
        self.selected_db = self.db_client[database]

    def select_collection(self, collection: str):
        self.selected_collection = self.selected_db[collection]

    def select_db_collection(self, database: str, collection: str):
        self.selected_collection = self.db_client[database][collection]

    def get_selected_collection(self):
        return self.selected_collection

    def insert_data(self, insert_datas):
        try:
            if len(insert_datas) == 1:
                insert_result = self.selected_collection.insert_one(insert_datas[0])
            else:
                insert_result = self.selected_collection.insert_many(insert_datas)
        except pymongo.errors.OperationFailure as e:
            return e.details
        return {"result": insert_result}

    def find_data_one(self, find_data, data_schema=None):
        if data_schema is None:
            data_schema = {}
        result = self.selected_collection.find_one(find_data, data_schema)
        return result

    def find_data_one_noid(self, find_data, data_schema=None):
        if data_schema is None:
            data_schema = {}
        data_schema = dict(data_schema, **{'_id': 0})
        result = self.selected_collection.find_one(find_data, data_schema)
        return result

    def find_data_many(self, find_data, data_schema=None):
        if data_schema is None:
            data_schema = {}
        result_list = list()
        results = self.selected_collection.find(find_data, data_schema)
        for result in results:
            result_list.append(result)
        return result_list

    def find_data_many_noid(self, find_data, data_schema=None):
        if data_schema is None:
            data_schema = {}
        data_schema = dict(data_schema, **{'_id': 0})
        result_list = list()
        results = self.selected_collection.find(find_data, data_schema)
        for result in results:
            result_list.append(result)
        return result_list


if __name__ == "__main__":
    db_client = MongoConnector()
    db_client.init_client()
    db_client.select_db_collection("trade", "dailystocks")
    insert_data = {
        "code": "601949.SH",
        "name": "中国出版",
        "trade_account": "17000000",
        "buy_qty": 1800,
        "sell_qty": 1800,
        "auto_trading": 1,
        "blacklist": 0,
        "first_trading": 1,
        "date": "2023-06-15",
        "timestamp": datetime.datetime.strptime("2023-06-15", "%Y-%m-%d")
    }
    result = db_client.insert_data([insert_data])
    if "errmsg" in result:
        print(result)
    result_id = str(result.inserted_id)
    # user = db_client.find_data_many_noid({})
    print(result_id)
    db_client.close_client()

